﻿#include "OnnxDetector.h"

#include <filesystem>

#include "inference.h"
#include "opencv2/opencv.hpp"

OnnxDetector::OnnxDetector() : detector_(std::make_shared<DCSP_CORE>()) {
  // 设置识别的类名
  detector_->classes = {
      "person",        "bicycle",      "car",
      "motorcycle",    "airplane",     "bus",
      "train",         "truck",        "boat",
      "traffic light", "fire hydrant", "stop sign",
      "parking meter", "bench",        "bird",
      "cat",           "dog",          "horse",
      "sheep",         "cow",          "elephant",
      "bear",          "zebra",        "giraffe",
      "backpack",      "umbrella",     "handbag",
      "tie",           "suitcase",     "frisbee",
      "skis",          "snowboard",    "sports ball",
      "kite",          "baseball bat", "baseball glove",
      "skateboard",    "surfboard",    "tennis racket",
      "bottle",        "wine glass",   "cup",
      "fork",          "knife",        "spoon",
      "bowl",          "banana",       "apple",
      "sandwich",      "orange",       "broccoli",
      "carrot",        "hot dog",      "pizza",
      "donut",         "cake",         "chair",
      "couch",         "potted plant", "bed",
      "dining table",  "toilet",       "tv",
      "laptop",        "mouse",        "remote",
      "keyboard",      "cell phone",   "microwave",
      "oven",          "toaster",      "sink",
      "refrigerator",  "book",         "clock",
      "vase",          "scissors",     "teddy bear",
      "hair drier",    "toothbrush",   "tank",
      "cross",
  };
  // 设置颜色盒子
  for (int i = 0; i < detector_->classes.size(); ++i) {
    // 随机颜色生成
    cv::RNG rng(cv::getTickCount());
    // 保存颜色
    colorBox_[i] = {rng.uniform(0, 256), rng.uniform(0, 256),
                    rng.uniform(0, 256)};
  }

  // 模型的路径
  const std::string model_path = "best.onnx";
  // 识别参数,使用CUDA识别
  DCSP_INIT_PARAM params{model_path, YOLO_ORIGIN_V8, {640, 640},
                         0.25,       0.45,           true};
  // 使用配置参数
  detector_->CreateSession(params);
}

OnnxDetector::~OnnxDetector() {
  cv::destroyAllWindows();
}

void OnnxDetector::detect(cv::Mat &inputImg) const {
  // 检测结果
  std::vector<DCSP_RESULT> res{};
  // 进行检测
  detector_->RunSession(inputImg, res);
  // 将检测结果进行绘制显示
  for (auto &re : res) {
    // 获得颜色
    cv::Scalar color(colorBox_.at(re.classId)[0], colorBox_.at(re.classId)[1],
                     colorBox_.at(re.classId)[2]);
    // 绘制矩形框
    cv::rectangle(inputImg, re.box, color, 3);

    // 置信度
    const float confidence = re.confidence;
    std::cout << std::fixed << std::setprecision(2);
    // 用于绘制的文字
    std::string label = detector_->classes[re.classId] + " "
        + std::to_string(confidence)
              .substr(0, std::to_string(confidence).size() - 4);

    // 绘制矩形框
    cv::rectangle(inputImg, cv::Point(re.box.x, re.box.y - 25),
                  cv::Point(re.box.x + label.length() * 15, re.box.y), color,
                  cv::FILLED);

    // 绘制文字
    cv::putText(inputImg, label, cv::Point(re.box.x, re.box.y - 5),
                cv::FONT_HERSHEY_SIMPLEX, 0.75, cv::Scalar(0, 0, 0), 2);
  }
}

void OnnxDetector::detectImages(const std::string &imgsFolder) {
  // 获得当前程序路径
  // std::filesystem::path current_path = std::filesystem::current_path();
  // 获得当前的图片文件夹路径
  // std::filesystem::path imgs_path = current_path / "images";
  // 遍历图片文件夹
  for (auto &i : std::filesystem::directory_iterator(imgsFolder)) {
    // 只识别图片文件
    if (i.path().extension() == ".jpg" || i.path().extension() == ".png"
        || i.path().extension() == ".jpeg") {
      // 读取图片
      std::string img_path = i.path().string();
      cv::Mat img = cv::imread(img_path);

      // 识别
      this->detect(img);

      // 显示识别的结果
      std::cout << "Press any key to next" << std::endl;
      cv::imshow("Result of Detection", img);
      cv::waitKey(0);
    }
  }
}

void OnnxDetector::detectVideo(const std::string &videoPath) {
  cv::VideoCapture capture{videoPath};
  // 一帧图片
  cv::Mat frame{};
  // 读取图片
  while (capture.isOpened() && capture.read(frame)) {
    // 识别
    this->detect(frame);

    // 显示识别的结果
    cv::imshow("Result of Detection", frame);
    // 如果按下了Q则退出(Q的键盘值是81)
    if (cv::waitKey(5) == 81) {
      break;
    }
  }
}
